import os
import cPickle
from itertools import izip,count
import numpy as np
import astro.spec as sp
import pylab as pl
from io import readtxt

# constants

Ckms = 299792.458         # speed of light km/s, exact
TOL = 20./Ckms            # tolerance to look for matching lines
TOLREST = 20./Ckms        # tolerance to look for matching rest wav lines
SIGNIF = 6.               # Significance in no. of EW sigma of a feature 
CUTOFF_V = 15000./Ckms    # rejection velocity red of lya emission 
ABSFILE = 'rdgen_ab_list' # name of rdgen abs list output file
REGION_V = 220./Ckms      # used writing out fitting regions
VDELT = 400./Ckms         # used plotting the the metals for confirmation
TRANSDIR = '/home/nhmc/programming/python/astro/'

def scanlines(specfile, **kwargs):
    """
    scanlines(spectrum_filename, [z=qso_redshift, vsig=vsigma,
              wa=list of waves])
    e.g
    scanlines('HE0940.fits', z=3.074, vsig=7.0)
    scanlines('HE0940.fits', wa=wavelengths, vsig=2.8)

    vsig in km/s.  Scans a list of wavelengths for redshifted metal
    transition groups (mostly doublets) given in atom.dat-style files
    in transitions/

    If it isn't given a list of wavelengths, it uses rdgen to generate
    a linelist redwards of the QSO lya emission line.  It then
    iterates over the wavelength list looking for lines at all
    redshifts matching positions of reference transitions.  Plots
    those found for confirmation.  Finally writes out confirmed
    systems, adding fitting regions and fitting with vpfit if
    requested.

    returns what?
    """

    # defining some constants
    ion_dict = dict(C_IV='c4.dat',   SiIV='si4.dat', MgII='mg2.dat',
                    FeII='fe2.dat', AlIII='al3.dat', CaII='ca2.dat',
                    SiII='si2.dat',   N_V='n5.dat',   NaI='na1.dat')

    spec = sp.read(specfile)
    spec.zqso = 99.0 if 'z' not in kwargs else kwargs['z']
    spec.vsig = 2.97 if 'vsig' not in kwargs else kwargs['vsig']
    wavmin = ( (1. + spec.zqso) * sp.W_LYA if 'wavmin' not in kwargs 
               else kwargs['wavmin'] )
    wavmax = spec.wa[-1] if 'wavmax' not in kwargs else kwargs['wavmax']

    print 'wavmin,wavmax are %s,%s' % (wavmin,wavmax)

    # skip these
    ## terrestrial molecular lines ('telluric'), identified
    ## from HE0940 spectrum
    ## 6276.95 - 6319.41
    ## 6867.17 - 7080.31
    ## 7155.89 - 7446.09
    ## 7593.19 - 7711.45
    ## 8103.37 - 8406.95?
    ## 8911.8 - 9865.42?

    # scan everything redwards of Lya emission with rdgen.  Generate
    # abline list.

    # rdgen input file.
    if 'wa' in kwargs:
        wav = kwargs['wa']    # already have list of wavelengths
    else:
        rdgen_input = "rd %s\nab %s\n\n\n\n%s,%s\nquit\n" % (
            specfile, ABSFILE, wavmin,wavmax)
        open('temporary_rdgen_input','w').write(rdgen_input)

        os.system('rdgen < temporary_rdgen_input > /dev/null')
        os.system('rm temporary_rdgen_input')

        # parse absorption line list

        # for now a simple-minded test - only wavelengths and ew and
        # errors of single absorption complexes (component positions in
        # comp, maybe use this info in the future).  Only want absorption
        # features, not emission.

        f = open(ABSFILE).readlines()
        del f[:7]
        ab = [];  comp = []
        for row in f:
            word = row.split(None,1)[0]
            if word == 'em':
                continue
            elif word == 'abs':
                ab.append(row)
                comp.append([])
            elif word == '***':
                comp[-1].append(row)

        wav, wav_er, ewidth, ewidth_er = readtxt(ab, usecols=(2,3,5,6))
        # quick look at lines
        pl.ion()
        spec.plot(draw=False)
        xmin,xmax,ymin,ymax = pl.axis()
        ticklen = (ymax - ymin) / 10.
        pl.ioff()
        for w in wav:
            i = spec.wa_index(w)
            pl.plot((w,w),(spec.co[i] + ticklen/10., spec.co[i] + ticklen),'k',
                    scaley=False,scalex=False,linewidth=1)
        pl.show()
        raw_input('Enter to continue ')

        significance = ew / ew_er
        # arrived at through trial and error - will depend on how accurate
        # error array is.
        condition = significance > SIGNIF
        wav = wav[condition]
        wav_er = wav_er[condition]
        ewidth = ewidth[condition]
        ewidth_er = ewidth_er[condition]

    # remove lines in areas contaminated by terrestrial molecular
    # absorption
    in_telluric = ((6276.95<wav)&(wav<6319.41) | (6867.17<wav)&(wav<7080.31) |
                   (7155.89<wav)&(wav<7446.09) | (7593.19<wav)&(wav<7711.45) |
                   (8103.37<wav)&(wav<8406.95) | (8911.8<wav) &(wav<9865.42) )
    wa = wav[~in_telluric]
##  wa_er = wav_er.compress(not_in_telluric)
##  ew = ewidth.compress(not_in_telluric)
##  ew_er = ewidth_er.compress(not_in_telluric)

    # Possible rest wavelength transitions:
    temp = raw_input('Scan for transitions?  ')
    if temp == 'y':
        sys = scan_all(ion_dict,spec.zqso,wa)
    temp = raw_input('Check transitions by eye?  ')
    if temp == 'y':
        good_sys = check(sys,spec)
    else:
        print 'Using systems from real_sys.cPickle'
        good_sys = cPickle.load(open('real_sys.cPickle'))

    # Have definite systems: Now see if there are any conflicting
    # assignments for single lines
    ids = []

    for i in xrange(len(wa)):
        matching_lines = []
        for ion_name in ion_dict.iterkeys():
            for s in good_sys[ion_name]:
                temp = [line for line in s if line[-2] == i]
                if len(temp) == 1:
                    matching_lines.append(temp[0])
                elif len(temp) > 2:
                    print ('More than one matching line found at a single '
                           'redshift.  Madness!  Madness, I tell you!')
        if len(matching_lines) == 0:
            ids.append(['unknown'])
        else:
            ids.append(matching_lines)

    multiple_ids = [(w,i) for w,i in izip(wa,ids) if len(i) > 1]
    print 'List of lines with multiple identifications:\n', multiple_ids
    print "Couldn't identify %s lines (total %s)" % (
        len([i for i in ids if i[0] =='unknown']), len(ids))

    # TODO: now use existing known transitions to search for others at
    # known zs. Later, only interested in C IV for now.

    temp = raw_input('Write out directories? ')
    if temp == 'y':
        writedir(good_sys,spec)

    if 'wa' not in kwargs:
        return good_sys,ids,wa,ab,comp
    else:
        return good_sys,ids

def scan_all(ion_dict,zqso,wa):
    """ Scans a spectrum for rest wavelength transitions, and then for
    redshifted QSO transitions.
    """
    # First rest transitions

    f = open('%stransitions/rest.dat' % TRANSDIR).readlines()
    f.sort()

    refnames = []
    refwavs = []
    refoscs = []

    for line in f:
        refnames.append(line[9:16])
        refwavs.append( float(line.split(None,1)[0]) )
        refoscs.append( float(line[16:].split(None,5)[4]) )

    sys = dict().fromkeys(ion_dict.keys())
    for key in sys:
        sys[key] = []

    # assign potential matches to each detected line

    rest_lines = []

    # look for z = 0 (Galactic,terrestrial) rest wavelength transitions
    # first.
    for refname,refwav,refosc in izip(refnames, refwavs, refoscs):
        j1 = wa.searchsorted(refwav * (1 - TOLREST))
        j2 = wa.searchsorted(refwav * (1 + TOLREST))
        for j,wav in izip(count(j1),wa[j1:j2]):
            line = (0.0, refname, refwav, refosc, j, wav)
            rest_lines.append(line)

    # Now search through the rest of the lines looking for redshifted
    # systems.

    for ion_name,ion_file in ion_dict.iteritems():
        l = scan(ion_file,wa)
        sys[ion_name].extend(l)
        # We can remove any systems that have z >> z_qso, as they
        # can't be real
        sys[ion_name] = [s for s in sys[ion_name] if
                         s[0][0] < zqso + CUTOFF_V * (1. + zqso)]

    # maybe add a loop here to join systems that are separated by less
    # than some velocity difference?

    return sys

def scan(ion_file,wa):
    """ ion_file = 'mg2.dat' or similar.  Scans spectrum for a series
    of redshifted qso transitions.  Returns a list of transitions.
    Finds the transitions from atom.dat style file in transitions/
    """
    f = open('%stransitions/%s' % (TRANSDIR,ion_file)).readlines()
    f.sort()
    refwavs = []
    refoscs = []
    systems = []
    for line in f:
        refwavs.append( float(line[6:].split(None,1)[0]) )
        refoscs.append( float(line[6:].split(None,2)[1]) ) # might use eventu.
    for i,wav in enumerate(wa):
        for j,refwav,refosc in izip(count(0),refwavs,refoscs):
            # assume wa[i] is redshifted instance of transition ref[j]
            zp1 = wav / refwav
            # find new expected transition wavelengths
            redshifted_refwav = np.array(refwavs) * zp1
            line = (zp1 - 1.0, refwav, refosc, i, wav)
            l = [line]
            # now see if any other transitions present
            for jj,zrefwav in enumerate(redshifted_refwav):
                if jj == j:  continue
                # find wavelength range in which a matching line is expected
                i1 = wa.searchsorted(zrefwav * (1. - TOL))
                i2 = wa.searchsorted(zrefwav * (1. + TOL))
                # take closest line in this range
                mindist = TOL
                for ii in xrange(i1,i2):
                    dist = abs(wa[ii] - zrefwav)/wa[ii]
                    if dist < mindist:
                        newline = (zp1 - 1.0, refwavs[jj],
                                   refoscs[jj], ii, wa[ii])
                        mindist = dist
                if mindist < TOL:  l.append(newline)
            # if 2 or more lines found, keep system
            if len(l) > 1:  systems.append(l)

    # remove identical systems that have been identified multiple
    # times.  This seems to be unnecessarily complicated!  Can't
    # see a neater way.

    ids = [ [(item[1],item[3]) for item in s] for s in systems ]
    for i in xrange(len(ids)):
        ids[i].sort()
        ids[i] = tuple(ids[i])   # needed to enable set(ids)
    for item in set(ids):
        if ids.count(item) > 1:
            while ids.count(item) > 1:
                index = ids.index(item)
                del systems[index]
                del ids[index]

    for s in systems:
        s.sort()
    systems.sort()
    return systems

def check(sys,spec):
    """ Check whether systems look real by eye.  Save those that do
    look real.
    """
    good_sys = dict().fromkeys(sys.keys())
    for key in good_sys:
        good_sys[key] = []
    print 'Accepted systems will be incrementally saved to good_sys.cPickle'

    # plot absorption systems for by eye confirmation/rejection
    # read in spectrum

    pl.ioff()
    pl.figure(figsize=(11,8))
    for ion_name in good_sys.iterkeys():
        numsys = len(sys[ion_name])
        if numsys == 0:  continue
        for j,s in enumerate(sys[ion_name]):
            # work out number of x and y plots
            numlines = len(s)
            if numlines < 5:
                nx = ny = 2
            else:
                nx = int(np.ceil(np.sqrt(numlines)))
                ny = int(np.ceil(numlines / float(nx)))
            for i,absline in enumerate(s):
                zp1 = absline[0] + 1.
                wrest = absline[1]
                osc = absline[2]
                comment = '%s  %6.4f  %7.2f  %4.2f' % (
                    ion_name,zp1-1.,wrest,osc)
                if len(s) == 2 and i == 1:
                    pl.subplot(int(ny),int(nx),3)
                else:  pl.subplot(int(ny),int(nx),i+1)
                xmin = wrest * zp1 * (1 - VDELT)
                xmax = wrest * zp1 * (1 + VDELT)
                spec.plot(limits=(xmin,xmax),ymax_fac=1.4)
                # plot tick mark showing position of line, along with comment
                pl.title(comment)
                ymin,ymax = pl.axis()[2:]
                tick_low = (ymax - ymin) * 0.75 + ymin
                tick_high = (ymax - ymin) * 0.95 + ymin
                pl.plot((wrest*zp1, wrest*zp1),(tick_low, tick_high),
                        'k', lw=2, scalex=False, scaley=False)
            pl.show()
            temp = raw_input(
                "%s: %s of %s; 'y' + return to keep, just enter to discard\n"
                % (ion_name, j+1, numsys))
            if temp == 'y':
                good_sys[ion_name].append(s)
                # save incrementally in case something goes wrong
                cPickle.dump(good_sys, open('good_sys.cPickle', 'w'))
            pl.ioff()
            pl.clf()
    return good_sys

def writedir(good_sys,spec):
    """ Writes out identified absorption systems to directories.
    Writes out fort.26 files and fitting regions, and runs vpfit if
    asked.
    """
    pl.close()
    pl.ioff()
    pl.figure(figsize=(11,8))
    for ion_type in good_sys.iterkeys():
        temp = raw_input('write dirs for ion %s? (y)' % ion_type)
        if temp.lower() == 'n':  continue
        if len(good_sys[ion_type]) == 0:
            print 'No systems for this ion, skipping'
            continue
        try:
            os.mkdir(ion_type)
        except OSError:
            char = raw_input('Directory %s already exists, overwrite? '
                             % ion_type)
            if char.lower()[0] == 'y':
                os.system('rm -rf %s' % ion_type)
                os.mkdir(ion_type)
            else:  continue
        os.chdir(ion_type)
        for system in good_sys[ion_type]:
            dirname = '%6.4f' % system[0][0]
            os.mkdir(dirname)
            os.chdir(dirname)
            #  f13 style file, fitting regions
            regions13 = []
            regions26 = []
            ions13 = []
            ions26 = []
            ion_name = []
            # plot first
            if len(system) < 5:
                nx = ny = 2
            else:
                nx = int(np.ceil(np.sqrt(len(system))))
                ny = int(np.ceil(len(system) / nx))
            for i,line in enumerate(system):
                if ion_type == 'CIV':
                    name = 'C IV'
                elif ion_type == 'NV':
                    name = 'N V'
                else:  name = ion_type.replace('_',' ')
                zp1 = line[0] + 1.0
                wrest = line[1]
                # guess fitting region
                minwav = zp1 * wrest * (1 - REGION_V)
                maxwav = zp1 * wrest * (1 + REGION_V)
                # plotting limits
                xmin = wrest * zp1 * (1 - 1.5 * REGION_V)
                xmax = wrest * zp1 * (1 + 1.5 * REGION_V)
                if len(system) == 2 and i+1 == 2:
                    pl.subplot('%s%s%s' % (nx,ny,3))
                else:
                    pl.subplot('%s%s%s' % (nx,ny,i+1))
                a = spec.plot(limits=(xmin,xmax),ymax_fac=1.4)
                ymin,ymax = pl.axis()[2:]
                tick_low = (ymax - ymin) * 0.75 + ymin
                tick_high = (ymax - ymin) * 0.95 + ymin
                # plot the feature position
                pl.plot((wrest*zp1,wrest*zp1),(tick_low, tick_high),
                        'k',lw=2,scalex=False,scaley=False)
                # and the guessed fitting region
                pl.plot((minwav,maxwav),(tick_high,tick_high),lw=3,
                        color='yellow',scalex=False,scaley=False)
            pl.show()
            # write out a line 
            for line in system:
                zp1 = line[0] + 1.0
                wrest = line[1]
                column = 13.0
                bpar = 2.0
                temp = raw_input(
                    'click new wavmin,wavmax or enter to continue: ')
                if len(temp) > 0:
                    minwav,maxwav = [float(item) for item in temp.split(',')]
                else:
                    minwav = zp1 * wrest * (1 - REGION_V)
                    maxwav = zp1 * wrest * (1 + REGION_V)
                regions26.append(' %%%% %s  1 %7.2f %7.2f vsig=%s\n' %
                               (spec.filename, minwav, maxwav, spec.vsig))
                regions13.append(' %s  1 %7.2f %7.2f vsig=%s\n' %
                               (spec.filename, minwav, maxwav, spec.vsig))
                if name in ion_name:  continue
                ions26.append(
                    ' %8s    %9.7f  0.0  %5.2f  0.0  %6.3f  0.0  0\n' %
                            (name, zp1 - 1., bpar, column))
                ions13.append(
                    ' %8s    %6.3f  %9.7f  %5.2f 0.00   0.00E+00  0\n' %
                            (name, column, zp1 - 1.0, bpar))
                ion_name.append(name)
            # fort.13 file for vpguess
            f13 = open('fort.13','w')
            f13.write('   *\n')
            f13.writelines(regions13)
            f13.write('  *\n')
            f13.writelines(ions13)
            f13.close()
            f26 = open('fort.26','w').writelines(regions26 + ions26)
            # run vpfit if asked
            temp = raw_input('Run vpfit? (n)')
            if temp == 'y':
                open('temporary_vpfit_input','w').write(
                    'f ad 3.0 il\n\nfort.26\nn\n\n')
                os.system('vpfit9 < temporary_vpfit_input')
                os.system('rm temporary_vpfit_input')
            os.chdir('..')
        os.chdir('..')
